package com.yww.datastructure.btree;

import java.util.Arrays;

/**
 * @author yww
 * @description 简单B树实现：平衡树
 * @since 2024/2/18 11:28
 */
public final class SimpleBTree {
    Node root; // 根节点
    final int t; // 节点的最小度
    final int MIN_KEY_NUMBER; // 最小有效key数
    final int MAX_KEY_NUMBER; // 最大有效key数

    public SimpleBTree() {
        this(2);
    }

    public SimpleBTree(int t) {
        this.t = t;
        MIN_KEY_NUMBER = t - 1;
        MAX_KEY_NUMBER = 2 * t - 1;
        root = new Node(t);
    }

    /**
     * 添加方法，将key放入b树
     *
     * @param key 要添加的key
     */
    public void put(int key) {
        doPut(null, 0, root, key);
    }

    /**
     * 递归添加
     *
     * @param parent    要添加key的节点的父节点
     * @param nodeIndex 要添加key的节点的索引
     * @param node      要添加key的节点
     * @param key       要往b树添加的key
     */
    private void doPut(Node parent, int nodeIndex, Node node, int key) {
        int i = 0;
        while (i < node.keyNumber) {
            if (node.keys[i] == key) { // 找到重复的key，走更新
                return; // 不涉及val，什么都不做
            } else if (node.keys[i] > key) { // 找到key要插入的位置
                break;
            }
            i++;
        }
        if (node.leaf) { // 如果节点是叶子节点，则直接进行插入
            node.insertKey(key, i);
        } else { // 如果节点是非叶子节点，则继续在children[i]处节点进行递归插入
            doPut(node, i, node.children[i], key);
        }
        if (node.keyNumber == MAX_KEY_NUMBER) { // 节点key数目达到满（进行限制），此时应该进行节点分裂
            split(parent, node, nodeIndex);
        }
    }

    /**
     * 分裂方法
     *
     * @param parent    要分裂节点的父节点
     * @param left      要进行分裂的节点（left是指，分离完成后该节点会插入到前一个左侧位置）
     * @param leftIndex 分裂节点在父节点的索引位置
     */
    protected void split(Node parent, Node left, int leftIndex) {
        if (parent == null) { // 表示要分裂的是根节点
            Node newRoot = new Node(t); // 需要创建新根
            newRoot.leaf = false;
            newRoot.insertChild(root, 0); // 旧根作为新根0孩子
            root = newRoot;
            parent = newRoot;
        }
        // 创建right节点（分裂后大于当前left节点）
        Node right = new Node(t);
        right.leaf = left.leaf; // 他们会在同一层
        // 把t以后的key和child都拷贝过去
        System.arraycopy(left.keys, t, right.keys, 0, t - 1);
        if (!left.leaf) {
            System.arraycopy(left.children, t, right.children, 0, t); // child会比key多一个
            for (int i = t; i < left.children.length; i++) {
                left.children[i] = null; // help gc
            }
        }
        right.keyNumber = t - 1; // 更新right的有效key数目
        left.keyNumber = t - 1; // 减小left的有效key数目
        // t-1处的key插入到parent的leftIndex处
        parent.insertKey(left.keys[t - 1], leftIndex);
        // right节点作为parent的孩子，插入到leftIndex+1处
        parent.insertChild(right, leftIndex + 1);
    }

    /**
     * 删除key
     */
    public void remove(int key) {
        doRemove(null, root, 0, key);
    }

    private void doRemove(Node parent, Node node, int index, int key) {
        int i = 0;
        while (i < node.keyNumber) {
            if (node.keys[i] >= key) {
                break;
            }
            i++;
        }
        // i 找到：代表删除key的索引
        // i 没找到：代表到第i个孩子继续查找
        if (node.leaf) { // 叶子节点
            if (!found(node, key, i)) { // case1 叶子节点没有找到
                return; // 直接结束
            } else { // case2 叶子节点找到
                node.removeKey(i); // 删除key
            }
        } else { // 非叶子节点
            if (!found(node, key, i)) { // case3 非叶子节点没有找到
                doRemove(node, node.children[i], i, key); // 让第i个孩子递归删除
            } else { // case4 非叶子节点找到
                // 第一步：找到后继key
                Node s = node.children[i + 1];
                while (!s.leaf) {
                    s = s.children[0]; // 往左边找
                }
                int sKey = s.keys[0];
                // 第二步：替换待删除的key
                node.keys[i] = sKey;
                // 第三步：递归删除后继节点的key
                doRemove(node, node.children[i + 1], i + 1, sKey);
            }
        }
        if (node.keyNumber < MIN_KEY_NUMBER) { // 节点key数目太少，需要进行平衡调整
            // 调整平衡 case 5、case 6
            balance(parent, node, index);
        }
    }

    /**
     * 对x节点进行平衡调整
     *
     * @param parent x的父节点
     * @param x      调整的x节点
     * @param i      x节点的索引
     */
    private void balance(Node parent, Node x, int i) {
        // case 6 要调整的是根节点
        if (x == root) {
            if (root.keyNumber == 0 && root.children[0] != null) {
                root = root.children[0];
            }
            return;
        }
        // 找到左兄弟和右兄弟
        Node left = parent.childLeftSibling(i);
        Node right = parent.childRightSibling(i);
        // case 5-1 左边富裕，进行右旋
        if (left != null && left.keyNumber > MIN_KEY_NUMBER) {
            // a) 父节点中前驱key旋转下来
            x.insertKey(parent.keys[i - 1], 0);
            // B) left中最大的孩子换爹
            if (!left.leaf) {
                x.insertChild(left.removeRightmostChild(), 0);
            }
            // c) left中最大的key旋转上去
            parent.keys[i - 1] = left.removeRightmostKey();
            return;
        }
        // case 5-2 右边富裕，进行左旋
        if (right != null && right.keyNumber > MIN_KEY_NUMBER) {
            // a) 父节点中后继key旋转下来
            x.insertKey(parent.keys[i], x.keyNumber);
            // b) right中最小的孩子换爹
            if (!right.leaf) {
                x.insertChild(right.removeLeftmostChild(), x.keyNumber);
            }
            // c) right中最小的key旋转上去
            parent.keys[i] = right.removeLeftmostKey();
            return;
        }
        // case 5-3 两边都不够借，向左合并
        if (left != null) { // 向左兄弟合并
            // 父节点删除节点
            parent.removeChild(i);
            left.insertKey(parent.removeKey(i - 1), left.keyNumber);
            x.moveToTarget(left);
        } else {
            // 向自己合并
            parent.removeChild(i + 1);
            x.insertKey(parent.removeKey(i), x.keyNumber);
            right.moveToTarget(x);
        }
    }

    private boolean found(Node node, int key, int i) {
        return i < node.keyNumber && node.keys[i] == key;
    }

    /**
     * 节点定义
     */
    static class Node {
        int[] keys; // 存储此节点的多个key，没有考虑value
        int keyNumber; // 当前节点的有效key数
        int t; // 当前节点的最小度数（决定了key的最小数t-1，最大数2t-1）
        Node[] children; // 当前节点的孩子
        boolean leaf = true; // 当前节点是否是叶子节点。默认true

        /**
         * @param t 设置当前节点的最小度
         */
        public Node(int t) {
            this.t = t;
            this.keys = new int[2 * t - 1];
            this.children = new Node[2 * t];
        }

        public Node(int[] keys) {
            this.keys = keys;
        }

        /**
         * 多路查找
         *
         * @param key
         */
        Node get(int key) {
            int i = 0;
            // 重小往大找
            while (i < keyNumber && keys[i] < key) {
                i++;
            }
            // 找到了
            if (i < keyNumber && keys[i] == key) {
                return this;
            }
            // 如果当前节点是叶子节点也没有找到，返回null
            if (leaf) {
                return null;
            }
            // 递归往子节点找，i即要查找的子节点
            return children[i].get(key);
        }

        /**
         * key插入方法
         *
         * @param key   要插入的key
         * @param index 插入的位置
         */
        void insertKey(int key, int index) {
            // 元素往后移
            System.arraycopy(keys, index, keys, index + 1, keyNumber - index);
            // 插入
            keys[index] = key;
            // 有效key数增加
            keyNumber++;
        }

        /**
         * 节点child插入方法
         *
         * @param child 要插入的孩子
         * @param index 插入的位置
         */
        void insertChild(Node child, int index) {
            // 元素往后移
            System.arraycopy(children, index, children, index + 1, keyNumber - index);
            // 插入
            children[index] = child;
        }

        // 移出指定index处的key
        int removeKey(int index) {
            int removed = keys[index];
            // index后元素往前移动
            // System.arraycopy(keys, index + 1, keys, index, (keyNumber - 1) - index);
            // keyNumber--;
            System.arraycopy(keys, index + 1, keys, index, --keyNumber - index);
            return removed;
        }

        // 移出最左边的key
        int removeLeftmostKey() {
            return removeKey(0);
        }

        // 移出最右边的key
        int removeRightmostKey() {
            return removeKey(keyNumber - 1);
        }

        // 移出index处的child
        Node removeChild(int index) {
            Node removed = children[index];
            System.arraycopy(children, index + 1, children, index, keyNumber - index);
            children[keyNumber] = null; // help gc
            return removed;
        }

        // 移出最左边的child
        Node removeLeftmostChild() {
            return removeChild(0);
        }

        // 移出最右边的child
        Node removeRightmostChild() {
            return removeChild(keyNumber);
        }

        // 获得index孩子左处的兄弟
        Node childLeftSibling(int index) {
            return index > 0 ? children[index - 1] : null;
        }

        // 获得index孩子右处的兄弟
        Node childRightSibling(int index) {
            return index == keyNumber ? null : children[index + 1];
        }

        // 以追加的方式复制当前节点的所有key和child到target
        void moveToTarget(Node target) {
            int start = target.keyNumber;
            // 复制所有child
            if (!leaf) {
                for (int i = 0; i <= keyNumber; i++) {
                    target.children[start + i] = children[i];
                }
            }
            // 复制所有key
            for (int i = 0; i < keyNumber; i++) {
                target.keys[target.keyNumber++] = keys[i];
            }
        }

        @Override
        public String toString() {
            return Arrays.toString(Arrays.copyOfRange(keys, 0, keyNumber));
        }
    }
}
